#include "stdafx.h"
#include "ProcessMgr.h"
#include "Encode.h"
#include "util.h"
#include "ps.h"
#include "DriverWrapper.h"
#include <set>

HICON GetIconFromFile(LPCTSTR szImageFileName);

CProcessMgr m_ProcessMgr;

CUniqueProcess::CUniqueProcess(DWORD dwProcessId, DWORD dwParentProcessId, FILETIME *createTime, bool bIs64Bit, DWORD dwSessionId, LPCWSTR szImagePath)
{
	m_bAlive = true;
	m_bIs64Bit = bIs64Bit;
	m_ProcessId = dwProcessId;
	m_ParentProcessId = dwParentProcessId;;
	m_CreateTime = *createTime;
	m_SessionId = dwSessionId;
	m_ImagePath = szImagePath;
	m_ProcessName = ExtractFileName(szImagePath);
	
	m_DisplayName = m_ProcessName;
	if (Is64Bit() && !m_bIs64Bit)
	{
		m_DisplayName += L" (*32)";
	}

	m_pParentProcess = NULL;
	m_TreeNode = NULL;
	m_Icon = NULL;
}

bool CUniqueProcess::operator==(const CUniqueProcess &key)const
{
	if (key.m_ProcessId == this->m_ProcessId && key.m_CreateTime.dwLowDateTime == this->m_CreateTime.dwLowDateTime && key.m_CreateTime.dwHighDateTime == this->m_CreateTime.dwHighDateTime)
		return true;

	return false;
}

CProcessMgr::CProcessMgr()
{
	InitializeCriticalSection(&m_Lock);
	m_hReadyEvent = CreateEvent(NULL, FALSE, FALSE, NULL);
}

CProcessMgr::~CProcessMgr()
{
	DeleteCriticalSection(&m_Lock);
	CloseHandle(m_hReadyEvent);
}

void CProcessMgr::Lock()
{
	EnterCriticalSection(&m_Lock);
}
void CProcessMgr::Unlock()
{
	LeaveCriticalSection(&m_Lock);
}

CUniqueProcess *CProcessMgr::Find(DWORD dwProcessId, FILETIME *CreateTime)
{
	if (m_KeyMap.empty())
		return NULL;

	CUniqueProcessKey uk(dwProcessId, CreateTime);

	CProcessKeyMap::iterator itor = m_KeyMap.find(uk);
	if (itor == m_KeyMap.end())
		return NULL;

	return itor->second;
}

CUniqueProcess *CProcessMgr::Find(DWORD dwProcessId)
{
	if (m_KeyMap.empty())
		return NULL;

	CProcessIdMap::iterator itor = m_IdMap.find(dwProcessId);
	if (itor == m_IdMap.end())
		return NULL;

	do 
	{
		if(itor->second->m_bAlive)
			return itor->second;
		itor++;
	} while (itor != m_IdMap.end());

	return NULL;
}

void CProcessMgr::InsertProcess(CUniqueProcess *up)
{
	m_List.push_back(up);

	CUniqueProcessKey upKey(up->m_ProcessId, &up->m_CreateTime);
	m_KeyMap[upKey] = up;
	m_IdMap.insert(std::pair<DWORD, CUniqueProcess *>(up->m_ProcessId, up));
}

void CProcessMgr::FillParent(CUniqueProcess *up)
{
	CUniqueProcess *parent = Find(up->m_ParentProcessId);
	if (parent)
	{
		up->m_pParentProcess = parent;
		parent->m_ChildProcesses.push_back(up);
	}
}

void CProcessMgr::FillIcon(CUniqueProcess *up)
{
	if (!up->m_ImagePath.empty())
	{
		HICON hIcon = NULL;
		CUniqueIconMap::iterator itor = m_UniqueIconMap.find(up->m_ImagePath);
		if (itor == m_UniqueIconMap.end())
		{
			hIcon = GetIconFromFile(up->m_ImagePath.c_str());
			up->m_Icon = new CUniqueIcon(hIcon, m_IconList.Add(hIcon));
			m_UniqueIconMap[up->m_ImagePath] = up->m_Icon;
		}
		else
		{
			up->m_Icon = itor->second;
		}
	}
}

void CProcessMgr::FillList(void)
{
	std::set<DWORD> v;

	NewEnumProcess(v);

	DWORD dwParentProcessId;
	FILETIME CreateTime;
	BOOLEAN bIs64Bit;
	DWORD SessionId;
	CString ImagePathNormalized;
	
	std::wstring wImagePath;
	TCHAR szImagePath[MAX_PATH] = { 0 };

	//Special: System Process...
	if (GetProcessCreateTime(4, &CreateTime))
	{
		if (GetNtoskrnlImagePath(szImagePath))
		{
			//replace SystemRoot
			if (!_wcsnicmp(szImagePath, L"\\SystemRoot\\", _ARRAYSIZE(L"\\SystemRoot\\") - 1))
			{
				TCHAR szSystemDirectory[MAX_PATH];
				GetSystemWindowsDirectory(szSystemDirectory, MAX_PATH);
				ImagePathNormalized.Format(L"%s\\%s", szSystemDirectory, szImagePath + _ARRAYSIZE(L"\\SystemRoot\\") - 1);
			}
			else
			{
				ImagePathNormalized = szImagePath;
			}

			//System process must be in Session 0
			CUniqueProcess *up = new CUniqueProcess(4, 0, &CreateTime, Is64Bit() ? true : false, 0, ImagePathNormalized);
			up->m_ProcessName = L"System";
			up->m_DisplayName = L"System";
			InsertProcess(up);
			FillIcon(up);
		}
	}

	for (std::set<DWORD>::iterator itor = v.begin(); itor != v.end(); ++itor)
	{
		DWORD dwProcessId = *itor;

		if (GetProcessBaseInfo(dwProcessId, &dwParentProcessId, &CreateTime, &bIs64Bit, &SessionId, wImagePath))
		{
			CUniqueProcess *up = new CUniqueProcess(dwProcessId, dwParentProcessId, &CreateTime, bIs64Bit ? true : false, SessionId, wImagePath.c_str());
			InsertProcess(up);
			FillIcon(up);
		}
	}

	for (size_t i = 0; i < m_List.size(); ++i)
	{
		CUniqueProcess *up = m_List[i];
		FillParent(up);
	}
}

void CProcessMgr::CallWndCallbacks(CUniqueProcess *up, UINT Msg)
{
	for (size_t i = 0; i < m_CallbackWnds.size(); ++i)
	{
		PostMessage(m_CallbackWnds[i], Msg, (WPARAM)up, 0);
	}
}

void CProcessMgr::RegisterWndCallbacks(HWND hWnd)
{
	m_CallbackWnds.push_back(hWnd);
}